"""
异步的方式下载视觉中国茶叶相关图片
https://www.vcg.com/creative-image/chaye/
"""
import asyncio
import hashlib
import json
import os
import re
import time
import tkinter
from tkinter import Tk

import aiofiles
import aiohttp
from lxml import etree


def md5(script_str):
    m = hashlib.md5()
    m.update(script_str.encode("utf8"))
    return m.hexdigest()


headers = {
    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/83.0.4103.61"
                  " Safari/537.36"
}


def parse_get_img_src(html):
    """
    解析到图片url
    :param html:列表页整体的html
    :return:图片url列表
    """
    ex = 'window.__PRELOADED_STATE__ = (.*?)</script><script'
    data = re.findall(ex, html)[0]
    json_obj = json.loads(data)
    list_data = json_obj['searchImage']['data']['list']
    flash_list = []
    for i in list_data:
        flash_data = {
            'unique_key': md5(str(i['id'])),
            'src': 'http:' + i['src'],
            'createdTime': i['createdTime'],
        }
        flash_list.append(flash_data)
    return flash_list


def get_dir():
    """ 不存在则创建目录"""
    # 项目跟路径
    root_path = './'
    img_dir = 'tea_images'
    directory = root_path + img_dir
    if not os.path.exists(directory):
        os.mkdir(directory)
    return directory


async def download_img(img_data):
    # 下载图片
    img_url = img_data['src']
    # print(img_url, '正在下载。。。')
    async with aiohttp.ClientSession() as session:
        async with await session.get(img_url) as response:
            b = await response.read()
            file_name = img_url.split('/')[-1]
            file_full_path = get_dir() + "/" + file_name
            await save_picture(file_full_path, b)


async def save_picture(file_full_path, b):
    # 如果文件存在，则跳过
    if os.path.exists(file_full_path):
        return
    # 异步方法写文件
    async with aiofiles.open(file_full_path, 'wb+') as fp:
        await fp.write(b)
        # print(file_full_path, '下载成功')


def save_data():
    # 保存到数据库
    pass


def upload_img_chr():
    # 上传图床
    pass


def parse_next_page(html):
    # 解析下一页
    selector = etree.HTML(html)
    next_page_path = "//div[@class='center_page']/a[@class='paginationNowPage']/following-sibling::a[1]/@href"
    # 下一页的url
    next_page_url_arr = selector.xpath(next_page_path)
    if next_page_url_arr:
        return next_page_url_arr[0]
    return None


class NextPageReader(object):
    """
    自定义异步迭代器，迭代读取下一页的列表页
    """

    def __init__(self, page):
        self.next_page = page
        self.detail_page_list = []

    async def get_list_page(self):
        """
        请求列表页的url，并解析，把下一页的url和当前列表页的url保存
        :return:
        """
        async with aiohttp.ClientSession() as session:
            async with await session.get(self.next_page, headers=headers) as response:
                # await+可等待的对象（协程对象、Future、Task对象） 是io等待
                print("请求url:%s,响应码是：%d" % (self.next_page, response.status))
                if response.status == 404:
                    return None
                else:
                    html = await response.text()
                    # 下一页的url
                    next_page_url = parse_next_page(html)
                    # 详情页数据
                    self.detail_page_list = parse_get_img_src(html)

                    if not next_page_url:
                        # 没有下一页
                        self.next_page = None
                    else:
                        self.next_page = next_page_url
                    return self

    def __aiter__(self):
        return self

    async def __anext__(self):
        # 如果下一页为空停止迭代
        if self.next_page is None:
            raise StopAsyncIteration

        val = await self.get_list_page()
        if val is None:
            raise StopAsyncIteration
        return val


async def main(url):
    # https://www.vcg.com/creative-image/chaye/
    # 列表页迭代器
    obj = NextPageReader(url)
    # 遍历时即启动迭代器，不断得到下一页的列表页
    async for item in obj:
        start = time.time()
        tasks = []
        for detail_pages in item.__getattribute__('detail_page_list'):
            # print(detail_pages)
            tasks.append(download_img(detail_pages))

        done, padding = await asyncio.wait(tasks)
        print('下载成功数' + str(len(done)))
        print("耗时：", time.time() - start)


# sy_url = input("输入要爬取的url：")
#
# print(sy_url)
# if sy_url:
#     asyncio.run(main(url=sy_url))
#
def call_back():
    sy_url = input_text.get()
    print(sy_url)
    if sy_url:
        asyncio.run(main(url=sy_url))


if __name__ == '__main__':
    root = Tk()
    root.title('视觉中国图片爬取工具')
    root.geometry('500x100+534+340')
    label = tkinter.Label(root, text="爬取链接")
    label.pack()

    v = tkinter.StringVar()
    v.set('https://www.vcg.com/creative-image/chaye/')
    input_text = tkinter.Entry(root, width=200, bd=5, textvariable=v)
    input_text.pack()

    submit_btn = tkinter.Button(root, text="确定", command=call_back)
    submit_btn.pack()
    # 进入消息循环
    root.mainloop()
